import jittor as jt
from jittor import nn
from model.resnet101_jittor import ResNetV1c
from model.cc_head import CCHead, FCNHead

class EncoderDecoder(nn.Module):
    def __init__(self, pretrained=None, resume=None, crit=None):
        super(EncoderDecoder, self).__init__()
        self.backbone = ResNetV1c(pretrained=pretrained)
        self.decode_head = CCHead(recurrence=2)
        self.auxiliary_head = FCNHead()
        if crit is None:
            self.crit = nn.CrossEntropyLoss(ignore_index=-1)
        else:
            self.crit = crit
        if resume is not None:
            self.load_parameters(jt.load(resume))
            # self.load(resume)
            print("resumed model loaded!")

    def pixel_acc(self, pred, label):
        preds = jt.argmax(pred, dim=1)[0].long()
        valid = (label >= 0).long()
        acc_sum = jt.sum(valid * (preds == label).long())
        pixel_sum = jt.sum(valid)
        acc = acc_sum.float() / (pixel_sum.float() + 1e-10)
        return acc

    def execute(self, feed_dict, *, segSize=None):
        img_data = jt.array(feed_dict['img_data'].numpy())
        seg_label = jt.array(feed_dict['seg_label'].numpy())
        # training
        if segSize is None:
            features = self.backbone(img_data)

            pred = self.decode_head(features[-1])
            pred_deepsup = self.auxiliary_head(features[-2])

            loss = self.crit(pred, seg_label)
            loss_deepsup = self.crit(pred_deepsup, seg_label)
            loss = loss + loss_deepsup * 0.4

            acc = self.pixel_acc(pred, seg_label)

            return loss, acc
        # inference
        else:
            pred = self.decode_head(self.backbone(img_data)[-1], segSize=segSize)
            return pred